#include "runge_kutta.h"
#include "start.h"

extern vofa_buff_t vofa;
IMU_t imu = {
    .quat = {.q0 = 1.0, .q1 = 0.0, .q2 = 0.0, .q3 = 0.0},
    .kp = 2.5,
    .ki = 0.001,
    .kd = 0};

// 计算四元数的模长
float quaternionLength(Quaternion_t q)
{
    return sqrt(q.q0 * q.q0 + q.q1 * q.q1 + q.q2 * q.q2 + q.q3 * q.q3);
}

// 归一化四元数
Quaternion_t normalizeQuaternion(Quaternion_t q)
{
    float len = quaternionLength(q);
    Quaternion_t result = {q.q0 / len, q.q1 / len, q.q2 / len, q.q3 / len};
    return result;
}

// 计算四元数乘法
Quaternion_t quaternionMultiply(Quaternion_t B, Quaternion_t W)
{
    Quaternion_t result;

    float Omega[4][4] = {
        {0, -W.q1 / 2, -W.q2 / 2, -W.q3 / 2},
        {W.q1 / 2, 0, W.q3 / 2, -W.q2 / 2},
        {W.q2 / 2, -W.q3 / 2, 0, W.q1 / 2},
        {W.q3 / 2, W.q2 / 2, -W.q1 / 2, 0},
    };

    /*更新矩阵*/
    result.q0 = B.q0 * Omega[0][0] + B.q1 * Omega[0][1] + B.q2 * Omega[0][2] + B.q3 * Omega[0][3];

    result.q1 = B.q0 * Omega[1][0] + B.q1 * Omega[1][1] + B.q2 * Omega[1][2] + B.q3 * Omega[1][3];

    result.q2 = B.q0 * Omega[2][0] + B.q1 * Omega[2][1] + B.q2 * Omega[2][2] + B.q3 * Omega[2][3];

    result.q3 = B.q0 * Omega[3][0] + B.q1 * Omega[3][1] + B.q2 * Omega[3][2] + B.q3 * Omega[3][3];

    return result;
}

/**
 * @brief 计算四元数的导数
 */
Quaternion_t quaternionDerivative(Quaternion_t q, float wx, float wy, float wz)
{
    Quaternion_t omega = {0, wx, wy, wz}; // 将角速度转换四元数
    return quaternionMultiply(q, omega);  // 计算导数
}

/** 龙格-库塔算法更新四元数
 * @param q 四元数
 * @param wx 角速度x分量
 * @param wy 角速度y分量
 * @param wz 角速度z分量
 * @param h 时间步长 单位：秒
 * @return 四元数
 */
Quaternion_t updateQuaternion(Quaternion_t q, float wx, float wy, float wz, float h)
{
    Quaternion_t k1, k2, k3, k4;

    // 计算 k1
    k1 = quaternionDerivative(q, wx, wy, wz);

    // 计算 k2
    Quaternion_t q2 = {q.q0 + h / 2 * k1.q0,
                       q.q1 + h / 2 * k1.q1,
                       q.q2 + h / 2 * k1.q2,
                       q.q3 + h / 2 * k1.q3};
    k2 = quaternionDerivative(q2, wx, wy, wz);

    // 计算 k3
    Quaternion_t q3 = {q.q0 + h / 2 * k2.q0,
                       q.q1 + h / 2 * k2.q1,
                       q.q2 + h / 2 * k2.q2,
                       q.q3 + h / 2 * k2.q3};
    k3 = quaternionDerivative(q3, wx, wy, wz);

    // 计算 k4
    Quaternion_t q4 = {q.q0 + h * k3.q0,
                       q.q1 + h * k3.q1,
                       q.q2 + h * k3.q2,
                       q.q3 + h * k3.q3};
    k4 = quaternionDerivative(q4, wx, wy, wz);

    // 更新四元数
    Quaternion_t qNew;
    qNew.q0 = q.q0 + h / 6 * (k1.q0 + 2 * k2.q0 + 2 * k3.q0 + k4.q0);
    qNew.q1 = q.q1 + h / 6 * (k1.q1 + 2 * k2.q1 + 2 * k3.q1 + k4.q1);
    qNew.q2 = q.q2 + h / 6 * (k1.q2 + 2 * k2.q2 + 2 * k3.q2 + k4.q2);
    qNew.q3 = q.q3 + h / 6 * (k1.q3 + 2 * k2.q3 + 2 * k3.q3 + k4.q3);

    // 归一化
    return normalizeQuaternion(qNew);
}

// 欧拉角转换为四元数函数 可用于初始状态矫正
// 初始欧拉角可用重力加速度计测得
// @param roll  绕x轴旋转角度 单位：弧度
// @param pitch 绕y轴旋转角度 单位：弧度
// @param yaw   绕z轴旋转角度 单位：弧度
Quaternion_t eulerToQuaternion(float roll, float pitch, float yaw)
{
    float cy = cos(yaw * 0.5);
    float sy = sin(yaw * 0.5);
    float cr = cos(roll * 0.5);
    float sr = sin(roll * 0.5);
    float cp = cos(pitch * 0.5);
    float sp = sin(pitch * 0.5);

    Quaternion_t q;
    q.q0 = cr * cp * cy + sr * sp * sy;
    q.q1 = sr * cp * cy - cr * sp * sy;
    q.q2 = cr * sp * cy + sr * cp * sy;
    q.q3 = cr * cp * sy - sr * sp * cy;

    return q;
}

/**
 *
 * @param quaternion 四元数
 * @param accx 加速度x分量
 * @param accy 加速度y分量
 * @param accz 加速度z分量
 * @param wx 角速度x分量
 * @param wy 角速度y分量
 * @param wz 角速度z分量
 * @param mx 磁力计x分量
 * @param my 磁力计y分量
 * @param mz 磁力计z分量
 * @param time 时间步长 单位：秒
 */
Quaternion_t updateAttitude(Quaternion_t quaternion, float ax, float ay, float az, float wx, float wy, float wz, float mx, float my, float mz, float time)
{

    /* 更新旋转矩阵 */
    /* 这两个转换矩阵 将b系  n系互相转换    */

    float cnb[3][3] = {
        {(imu.quat.q0 * imu.quat.q0 + imu.quat.q1 * imu.quat.q1 - imu.quat.q2 * imu.quat.q2 - imu.quat.q3 * imu.quat.q3), 2 * (imu.quat.q1 * imu.quat.q2 + imu.quat.q0 * imu.quat.q3), 2 * (imu.quat.q1 * imu.quat.q3 - imu.quat.q0 * imu.quat.q2)},
        {2 * (imu.quat.q1 * imu.quat.q2 - imu.quat.q0 * imu.quat.q3), imu.quat.q0 * imu.quat.q0 - imu.quat.q1 * imu.quat.q1 + imu.quat.q2 * imu.quat.q2 - imu.quat.q3 * imu.quat.q3, 2 * (imu.quat.q2 * imu.quat.q3 + imu.quat.q0 * imu.quat.q1)},
        {2 * (imu.quat.q1 * imu.quat.q3 + imu.quat.q0 * imu.quat.q2), 2 * (imu.quat.q2 * imu.quat.q3 - imu.quat.q0 * imu.quat.q1), imu.quat.q0 * imu.quat.q0 - imu.quat.q1 * imu.quat.q1 - imu.quat.q2 * imu.quat.q2 + imu.quat.q3 * imu.quat.q3}};

    /* cbn矩阵是cnb矩阵的转置矩阵 */
    float cbn[3][3] = {
        {cnb[0][0], cnb[1][0], cnb[2][0]},
        {cnb[0][1], cnb[1][1], cnb[2][1]},
        {cnb[0][2], cnb[1][2], cnb[2][2]}};

    // 加速度计 绕y轴旋转35° sin35=0.5877852522924731 cos35=0.8090169943749475
    float sin_i = 0;
    float cos_i = 1;
    /* 安装 旋转矩阵 (针对Y轴的旋转) -> CWW为正方向 */
    float Axis_matrix[3][3] = {
        {cos_i, 0, sin_i},
        {0, 1, 0},
        {-sin_i, 0, cos_i},
    };

    /* 数据输入 */
    imu.acc.x = ax;
    imu.acc.y = ay;
    imu.acc.z = az;
    imu.gyro.x = wx;
    imu.gyro.y = wy;
    imu.gyro.z = wz;

    imu.mag.x = mx ;
    imu.mag.y = my ;
    imu.mag.z = mz ;

    /* ****************预处理环节**************** */
    /* 磁力计安装坐标旋转 */
    imu.mag.x = imu.mag.x * Axis_matrix[0][0] + imu.mag.y * Axis_matrix[1][0] + imu.mag.z * Axis_matrix[2][0];
    imu.mag.y = imu.mag.x * Axis_matrix[0][1] + imu.mag.y * Axis_matrix[1][1] + imu.mag.z * Axis_matrix[2][1];
    imu.mag.z = imu.mag.x * Axis_matrix[0][2] + imu.mag.y * Axis_matrix[1][2] + imu.mag.z * Axis_matrix[2][2];

    /*磁力计归一化*/ // TODO: 磁力计与imu的坐标系纠正
    imu.mag = normalizeacceleration(imu.mag);

    /* 磁力计坐标转换到当前坐标系 cbn*mag  生成是世界坐标下的磁力计投影 */
    /* cbn */
    float hx = (cbn[0][0] * imu.mag.x + cbn[0][1] * imu.mag.y + cbn[0][2] * imu.mag.z); // hx
    float hy = (cbn[1][0] * imu.mag.x + cbn[1][1] * imu.mag.y + cbn[1][2] * imu.mag.z); // hy
    float hz = (cbn[2][0] * imu.mag.x + cbn[2][1] * imu.mag.y + cbn[2][2] * imu.mag.z); // hz

    /* ==将磁力计向量扭转 一直指向x轴 完成数据更新== */
    /* 获取上一次的磁力计向量修正 */
    float bx= imu.mag_b.x;
    float by= imu.mag_b.y;
    float bz= imu.mag_b.z;

    imu.mag_b.x  = sqrt(hx * hx + hy * hy);
    imu.mag_b.y = 0;
    imu.mag_b.z = hz;

    /*通过地磁计算角速度 西北天*/

    float roll_rate, pitch_rate, yaw_rate=0;
    roll_rate= atan2f(imu.mag.y, imu.mag.z);
    pitch_rate= atan2f(imu.mag.x, imu.mag.z) ;
    if(imu.mag.y!=0)
    {
    yaw_rate= atan2f(imu.mag.x, imu.mag.y) ;//一四象限
    }       
    else if(imu.mag.y==0 )
    {
        if(imu.mag.x>0)
        {
            yaw_rate= 3.1415926/2;
        }
        else if(imu.mag.x<0)
        {
            yaw_rate= -3.1415926/2;
        }
        
    }
    static float las_z=0;
    static float las_d=0;
    if (yaw_rate<-1.57&&las_z>1.57)
    {
        imu.mag_diff.z=3.1415926+yaw_rate-las_z;
    }
    if(yaw_rate>1.57&&las_z<-1.57)
    {
        imu.mag_diff.z=-3.1415926+yaw_rate-las_z;
    }
    
    imu.mag_diff.z=las_d + 0.2*(((yaw_rate-las_z)/time)-las_d);

    las_d=imu.mag_diff.z;
    las_z=yaw_rate;
    vofa.buff[7]=imu.mag_diff.z;
    vofa.buff[8]=imu.gyro.z;
    vofa.buff[9]=mix(imu.mag_diff.z,imu.gyro.z);

    imu.gyro.z =mix(imu.mag_diff.z,imu.gyro.z);


    /* 加速度归一化 */
    imu.acc = normalizeacceleration(imu.acc);

    float acc_vector[3] = {imu.acc.x, imu.acc.y, imu.acc.z};        // 加速度向量缓冲
    float gyro_vector[3] = {imu.gyro.x, imu.gyro.y, imu.gyro.z}; // 角速度向量缓冲

    /* 加速度安装坐标旋转*/
    imu.acc.x = acc_vector[0] * Axis_matrix[0][0] + acc_vector[1] * Axis_matrix[1][0] + acc_vector[2] * Axis_matrix[2][0];
    imu.acc.y = acc_vector[0] * Axis_matrix[0][1] + acc_vector[1] * Axis_matrix[1][1] + acc_vector[2] * Axis_matrix[2][1];
    imu.acc.z = acc_vector[0] * Axis_matrix[0][2] + acc_vector[1] * Axis_matrix[1][2] + acc_vector[2] * Axis_matrix[2][2];

    /* 角速度安装坐标旋转*/
    imu.gyro.x = gyro_vector[0] * Axis_matrix[0][0] + gyro_vector[1] * Axis_matrix[1][0] + gyro_vector[2] * Axis_matrix[2][0];
    imu.gyro.y = gyro_vector[0] * Axis_matrix[0][1] + gyro_vector[1] * Axis_matrix[1][1] + gyro_vector[2] * Axis_matrix[2][1];
    imu.gyro.z = gyro_vector[0] * Axis_matrix[0][2] + gyro_vector[1] * Axis_matrix[1][2] + gyro_vector[2] * Axis_matrix[2][2];

#if (MAG_correct)
    /* 获得矫正向量 */
    /* [bx,by,bz] 是磁力计在世界坐标系下的向量 只不过by被归为0 我们希望地磁极永远指向北 */
    float mag_wx = (cnb[0][0] * bx + cnb[0][1] * by + cnb[0][2] * bz); // 转换为B系下坐标系 地磁计的预计值
    float mag_wy = (cnb[1][0] * bx + cnb[1][1] * by + cnb[1][2] * bz); //
    float mag_wz = (cnb[2][0] * bx + cnb[2][1] * by + cnb[2][2] * bz); //
#endif

#if (ACC_correct)

    /* 四元数在B系下坐标系下的分量 */
    /*
    cnb * [0,0,1]
    [0,0,1]就是世界坐标系下的重力向量
    */
    float vec_x = cnb[0][2] * 0.5;
    float vec_y = cnb[1][2] * 0.5;
    float vec_z = cnb[2][2] * 0.5;
#endif


/* ***********************  姿态修正环节  ************************** */
#if (ACC_correct && MAG_correct)

    imu.error.x = (imu.acc.y * vec_z - imu.acc.z * vec_y) + (imu.mag.y * mag_wz - imu.mag.z * mag_wy);
    imu.error.y = (imu.acc.z * vec_x - imu.acc.x * vec_z) + (imu.mag.z * mag_wx - imu.mag.x * mag_wz);
    imu.error.z = (imu.acc.x * vec_y - imu.acc.y * vec_x) + (imu.mag.x * mag_wy - imu.mag.y * mag_wx);
#endif

#if (ACC_correct && !MAG_correct)
    {
        imu.error.x = (imu.acc.y * vec_z - imu.acc.z * vec_y);
        imu.error.y = (imu.acc.z * vec_x - imu.acc.x * vec_z);
        imu.error.z = (imu.acc.x * vec_y - imu.acc.y * vec_x);
    }
#endif

#if (!ACC_correct && MAG_correct)
    {
        imu.error.x = (imu.mag.y * mag_wz - imu.mag.z * mag_wy);
        imu.error.y = (imu.mag.z * mag_wx - imu.mag.x * mag_wz);
        imu.error.z = (imu.mag.x * mag_wy - imu.mag.y * mag_wx);
    }
#endif

#if (!ACC_correct && !MAG_correct)
    {
        imu.error.x = 0;
        imu.error.y = 0;
        imu.error.z = 0;
    }
#endif

    /* 更新误差微分 */
    imu.error_differ.x = imu.error.x - imu.error_last.x;
    imu.error_differ.y = imu.error.y - imu.error_last.y;
    imu.error_differ.z = imu.error.z - imu.error_last.z;
    /* 更新积分误差 */
    imu.error_integ.x += imu.error.x;
    imu.error_integ.y += imu.error.y;
    imu.error_integ.z += imu.error.z;

    /* 叠加角速度修正量 */
    imu.gyro.x = imu.gyro.x + imu.kp * imu.error.x + imu.error_integ.x * imu.ki + imu.error_differ.x * imu.kd;
    imu.gyro.y = imu.gyro.y + imu.kp * imu.error.y + imu.error_integ.y * imu.ki + imu.error_differ.y * imu.kd;
    imu.gyro.z = imu.gyro.z + imu.kp * imu.error.z + imu.error_integ.z * imu.ki + imu.error_differ.z * imu.kd;

    imu.error_last.x = imu.error.x; // 更新上一次的误差
    imu.error_last.y = imu.error.y;
    imu.error_last.z = imu.error.z;

    return updateQuaternion(quaternion, imu.gyro.x, imu.gyro.y, imu.gyro.z, time); // 更新四元数 输入数据为B系下数据
}

/*向量归一化*/
Vector3f_t normalizeacceleration(Vector3f_t vector)
{

    float len = sqrt(vector.x * vector.x + vector.y * vector.y + vector.z * vector.z);
    Vector3f_t result = {vector.x / len, vector.y / len, vector.z / len};
    return result;
}
